/*
 * HotPixelFinder.cpp
 *
 *  Created on: Jan 16, 2014
 *      Author: tbergmueller
 */

#include "HotPixelFinder.h"
#include <iostream>
#include <vector>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/core/core.hpp>
#include "utilities.h"


using namespace std;
using namespace cv;



/**
 * Calculates the mean value for each pixel in a series of images
 * @param inputImages series of images
 * @return Mean-matrix as a CV_32FC1 (float precision)
 */
cv::Mat calcPixelMean(const vector<cv::Mat>& inputImages)
{
	assert(inputImages.size() > 2);
	Mat mean = Mat::zeros(inputImages[0].size(), CV_32FC1);

	float fact = (float)(1.0/inputImages.size());

	for(vector<cv::Mat>::const_iterator it = inputImages.begin(); it != inputImages.end(); it++)
	{
		assert(it->size() == mean.size());
		for(int x=0; x<it->cols; x++)
		{
			for(int y=0; y<it->rows; y++)
			{
				mean.at<float>(y,x) += (float)(it->at<uchar>(y,x)) *fact ;
			}
		}
	}

	return mean;
}


/**
 * Computes the defect matrix by computing the mean value (which is output to oMean), which is then blurred and subtracted
 * from the unblurred mean value.
 *
 * Note this works in equally distributed grayvalues (random image content) only, thus correlated regions of images have to be masked out
 * afterwards.
 *
 * The difference between the mean and blurred mean is 0 if all pixels are the same and positive, if some pixel stings out, because it has
 * a certain offset that was added to each pixel in the original images (=hot)
 *
 * @param inputImages Vector of input images
 * @param oMean Output of the calculated mean-grayscale-image
 * @param medianKernelSize Kernel size of the median for medianBlur-method
 * @return Defect Mat, containing the difference between the mean and the blurred mean image.
 */
Mat getDefectMat(const vector<Mat>& inputImages, Mat& oMean, int medianKernelSize = 3)
{
	Mat mean = calcPixelMean(inputImages);
	oMean = mean;
	Mat mean_illus(mean.size(), CV_8UC1);

	for(int x=0; x<mean.cols; x++)
	{
		for(int y=0; y<mean.rows; y++)
		{
			mean_illus.at<uchar>(y,x) = (uchar) mean.at<float>(y,x);
		}
	}

	//imshow("mean", mean_illus);
	//waitKey();
	Mat smoothed;
	medianBlur(mean_illus, smoothed, medianKernelSize);

	Mat defects;
	subtract(mean_illus,smoothed,defects);

	return defects;
}




void writeImg(string prefix, string name, Mat& toWrite)
{
	stringstream ss;
	ss << prefix << "_" << name << ".png";

	//cout << ss.str() << endl;
	imwrite(ss.str(),toWrite);
}

/**
 * Counts the defects in an image. A defect is a pixel that is not zero
 * @param defectLocationMat Locations of the defects
 * @return Nr. of defective pixels
 */
int countDefects(const Mat& defectLocationMat)
{
	int cnt = 0;

	for(int x=0; x<defectLocationMat.cols; x++)
	{
		for(int y=0; y<defectLocationMat.rows; y++)
		{
			if(defectLocationMat.at<uchar>(y,x) != 0)
			{
				cnt++;
			}
		}
	}
	return cnt;
}

/**
 * Computes the defect-locations of some inputs.
 *
 * This is done by first calculating the Defect-Matrix (output via o_defects) and then applying a threshold on this
 * matrix. Each pixel, whose offset from its neighbourhood exceeds a certain threshold is classified to be hot
 * These locations are then output via o_defec_locations.
 *
 * By counting the defect-locations of two different aged image sets, the growth rate of hot pixel defects can be calculated
 *
 * @param inputs Input images, loaded with loadBitmaps()-method in utilities.h
 * @param output_file_prefix Prefix (Folder) of the files to output. Output is for debugging only
 * @param o_defects Output matrix of defect-matrix calculated with getDefectMat()
 * @param o_defect_locations matrix of defect-locations.
 * @param oMean output of mean grayscale image of inputs
 */
void HotPixelFinder::calcHotPixelLocations(vector<Mat>& inputs, string output_file_prefix, Mat& o_defects, Mat& o_defect_locations, Mat& oMean)
{
	cout << "Calculate defects for " << output_file_prefix << endl;

	o_defects = getDefectMat(inputs, oMean, MEDIAN_KERNEL_SIZE);


	Mat d; // defect working mat for classification
	normalize(o_defects,d,0,255, NORM_MINMAX, CV_8UC1);
	writeImg(output_file_prefix, "defects_with_correlated_data",d);

	// Do classification by thresholding, FIXME: Hardcoded parameters
	threshold(d,d,65,255,CV_THRESH_BINARY);
	writeImg(output_file_prefix, "defects_classified_with_correlated_data",d);

	// Remove correlated region (replace with black)
	rectangle(d,Point(0,0), Point(d.cols-1, 50),Scalar::all(0),CV_FILLED);
	rectangle(d,Point(100,50), Point(d.cols-100, d.rows-50),Scalar::all(0),CV_FILLED);
	writeImg(output_file_prefix, "defects_classified",d);


	o_defect_locations = d;

}

/**
 * Computes the amplitude of the hot pixels. This is done by subtracting the mean of gray-values in a hot pixels surrounding
 * region.
 * @param values
 * @param medianKernelSize
 */
void HotPixelFinder::calcHotPixelAmplitude(Mat& locations, Mat& values, int medianKernelSize)
{
	_hotPixelAmplitudes.clear();

	assert(medianKernelSize % 2 == 1); // has to be odd

	int k = (int)medianKernelSize / 2;
	for(int x=0; x<locations.cols; x++)
	{
		for(int y=0; y<locations.rows; y++)
		{
			if(locations.at<uchar>(y,x) != 0)
			{

				if( x-k < 0 || y-k < 0 || x+k > locations.cols-1 || y+k > locations.rows-1)
				{
					continue;
				}
				Mat a = values(Rect(x-k,y-k,medianKernelSize,medianKernelSize));

				float mean = 0;

				for(int ix=0; ix<a.cols; ix++)
				{
					for(int iy=0; iy<a.rows; iy++)
					{
						mean += a.at<float>(iy,ix);
					}
				}

				mean /= medianKernelSize*medianKernelSize;

				float diff = values.at<float>(y,x) - mean;
				//cout << diff << endl;
				_hotPixelAmplitudes.push_back(diff);

			}
		}
	}
}


void HotPixelFinder::readHotPixelAmplitudesTo(float oMean, float& oMin, float& oMax)
{
	 oMin = 1000000.0;
	 oMax = -1.0;

	 oMean = 0;

	for(int i=0; i<(int)_hotPixelAmplitudes.size(); i++)
	{
		oMean += _hotPixelAmplitudes[i];

		if(_hotPixelAmplitudes[i] < oMin)
			oMin = _hotPixelAmplitudes[i];

		if(_hotPixelAmplitudes[i] > oMax)
			oMax = _hotPixelAmplitudes[i];
	}

	oMean /= (float)_hotPixelAmplitudes.size();
}


void HotPixelFinder::process()
{
	vector<Mat> inputs;

	Mat defects_2009, defect_loc_2009;
	Mat mean_2009, mean_2013;
	Mat defects_2013, defect_loc_2013;
	Mat matches;


	loadBitmaps(_folder2009, _nrOfFiles, inputs);
	calcHotPixelLocations(inputs,"img/2009",defects_2009, defect_loc_2009, mean_2009);

	loadBitmaps(_folder2013, _nrOfFiles, inputs);
	calcHotPixelLocations(inputs,"img/2013",defects_2013, defect_loc_2013, mean_2013);

	bitwise_and(defect_loc_2009, defect_loc_2013, matches);

	imshow("Defect Locations 2009", defect_loc_2009);
	imshow("Defect Locations 2013", defect_loc_2013);
	imshow("Defect Locations 2009 that match 2013", matches);

	cout << "Defects in 2009 are " << countDefects(defect_loc_2009) << endl;
	cout << "Defects in 2013 are " << countDefects(defect_loc_2013) << endl;
	cout << "Defects that match are " << countDefects(matches) << endl;

	if(countDefects(defect_loc_2009))
	{
		double correctnessRatio = (double)countDefects(matches) / (double)countDefects(defect_loc_2009);
		cout << "Correctness ratio is " << correctnessRatio << endl;
	}

	// TODO: this function has to be called twice at the moment, once with mean_2013 and once with mean_2009. Make some better interface
	calcHotPixelAmplitude(matches,mean_2013, MEDIAN_KERNEL_SIZE);


	waitKey();

}




HotPixelFinder::HotPixelFinder(std::string folder_2009_data, std::string folder_2013_data, int nrOfImages)
{

	_nrOfFiles = nrOfImages;
	_folder2009 = folder_2009_data;
	_folder2013 = folder_2013_data;
}


HotPixelFinder::~HotPixelFinder()
{

}

